"""
清空所有文档记录和向量文件
"""
import sys
import os

# 添加项目根目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
backend_dir = os.path.dirname(current_dir)
sys.path.insert(0, backend_dir)

from sqlalchemy.orm import Session
from app.models.database import engine
from app.models.document import Document
from app.core.config import settings
import logging
import shutil

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def clear_all_documents():
    """清空所有文档记录和相关文件"""
    db = Session(engine)
    try:
        # 1. 查询所有文档
        documents = db.query(Document).all()
        
        if not documents:
            logger.info("数据库中没有文档记录")
        else:
            logger.info(f"找到 {len(documents)} 个文档记录")
            
            # 2. 删除物理文件
            deleted_files = 0
            for doc in documents:
                try:
                    # 删除文档文件
                    if doc.file_path and os.path.exists(doc.file_path):
                        os.remove(doc.file_path)
                        deleted_files += 1
                        logger.info(f"  已删除文件: {doc.file_path}")
                    
                    # 删除向量文件
                    if doc.vector_store_id:
                        vector_path = os.path.join(settings.VECTOR_STORE_PATH, f"{doc.vector_store_id}.faiss")
                        if os.path.exists(vector_path):
                            os.remove(vector_path)
                            logger.info(f"  已删除向量: {vector_path}")
                except Exception as e:
                    logger.warning(f"  删除文件失败: {str(e)}")
            
            logger.info(f"已删除 {deleted_files} 个物理文件")
            
            # 3. 删除数据库记录
            db.query(Document).delete()
            db.commit()
            logger.info(f"已删除 {len(documents)} 条数据库记录")
        
        # 4. 清理上传目录（可选）
        upload_path = settings.UPLOAD_PATH
        if os.path.exists(upload_path):
            file_count = len([f for f in os.listdir(upload_path) if os.path.isfile(os.path.join(upload_path, f))])
            if file_count > 0:
                logger.info(f"\n上传目录中还有 {file_count} 个文件")
                clean_upload = input("是否清空上传目录？(yes/no): ")
                if clean_upload.lower() in ['yes', 'y']:
                    for filename in os.listdir(upload_path):
                        file_path = os.path.join(upload_path, filename)
                        try:
                            if os.path.isfile(file_path):
                                os.remove(file_path)
                            elif os.path.isdir(file_path):
                                shutil.rmtree(file_path)
                        except Exception as e:
                            logger.error(f"删除 {file_path} 失败: {str(e)}")
                    logger.info("已清空上传目录")
        
        # 5. 清理向量存储目录（可选）
        vector_path = settings.VECTOR_STORE_PATH
        if os.path.exists(vector_path):
            vector_count = len([f for f in os.listdir(vector_path) if f.endswith('.faiss')])
            if vector_count > 0:
                logger.info(f"\n向量存储目录中还有 {vector_count} 个向量文件")
                clean_vectors = input("是否清空向量存储目录？(yes/no): ")
                if clean_vectors.lower() in ['yes', 'y']:
                    for filename in os.listdir(vector_path):
                        if filename.endswith('.faiss'):
                            file_path = os.path.join(vector_path, filename)
                            try:
                                os.remove(file_path)
                            except Exception as e:
                                logger.error(f"删除 {file_path} 失败: {str(e)}")
                    logger.info("已清空向量存储目录")
        
        logger.info("\n" + "="*60)
        logger.info("清理完成！")
        logger.info("="*60)
        
    except Exception as e:
        logger.error(f"清理过程出错: {str(e)}", exc_info=True)
        db.rollback()
    finally:
        db.close()

if __name__ == "__main__":
    import sys
    # 设置 UTF-8 输出
    if sys.platform == 'win32':
        import io
        sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    
    print("="*60)
    print("警告：此操作将删除所有文档记录和文件！")
    print("="*60)
    print("将要执行以下操作：")
    print("  1. 删除数据库中的所有文档记录")
    print("  2. 删除所有上传的文档文件")
    print("  3. 删除所有向量索引文件")
    print("  4. 可选：清空上传目录")
    print("  5. 可选：清空向量存储目录")
    print("="*60)
    
    confirm = input("\n确认继续？(yes/no): ")
    if confirm.lower() in ['yes', 'y']:
        clear_all_documents()
    else:
        print("已取消")

